import re
from dalchemy.parsers import DalchemyParser
# from dalchemy.parsers import register_parser
from dalchemy.data import TextHelper

txthelper = TextHelper()

__all__ = ["SingleLineParser",
           "QALinesParser",
           "LinesParser",
           "TextBlocksParser",
           ]

# @register_parser("line")
class SingleLineParser(DalchemyParser):
    """Parse the output of an LLM
    """
    prefix: str = "answer"

    def parse(self, text: str):
        """Parse the output to qa pairs."""
        line_regex = re.compile(rf'{self.prefix}[\d]*[:：]\s*(?P<line>.*?)$')  # 单行回答直接匹配到末尾
        m_line = line_regex.search(text)
        line = ""
        if m_line:
            line = m_line.group('line').strip()
        return [{self.prefix: line}]

    def get_format_instructions(self):
        format_str = f'''
        请以单行输出，格式如下：
        {self.prefix}: 回答。
        '''
        return format_str

# @register_parser("lines")
class LinesParser(DalchemyParser):
    """ 解析多行输出，每行具有相同的前缀"prefix:"，并以\n结束。用于批量生成多个问题或回答。
        Example:
            from dalchemy.parsers import LinesParser
            input_str = '''
                question: AI的发展历史是什么？它是如何从概念演变为现实应用的？
                question: AI的不同类型有哪些？它们之间的区别是什么？
                question: AI在医疗领域有哪些应用？它如何改善医疗保健和诊断过程？
            '''
            parser = LinesParser(prefix = "question" )
            res = parser.parse(input_str)
            print(res)
            >> [{'question': 'AI的发展历史是什么？它是如何从概念演变为现实应用的？'},
                {'question': 'AI的不同类型有哪些？它们之间的区别是什么？'},
                {'question': 'AI在医疗领域有哪些应用？它如何改善医疗保健和诊断过程？'}]
    """
    prefix: str = "answer"

    def parse(self, text: str):
        """Parse the output to qa pairs."""
        answer_regex = re.compile(rf'{self.prefix}[\d]*[:：]\s*(?P<ans>.*?)\n')
        m_a = re.finditer(answer_regex, text)
        answers = []
        if m_a:
            answers = [a['ans'].strip() for a in m_a]
        # assert len(questions) == len(answers), "question and answer 's length mast same!"
        res = []
        for ans in answers:
            res.append({self.prefix: ans})
        return res

    def get_format_instructions(self):
        format_str = f'''
        请以下面的格式输出：
        {self.prefix}: 回答1\n
        {self.prefix}: 回答2\n
        ...
        '''
        return format_str

# @register_parser("qa_lines")
class QALinesParser(DalchemyParser):
    """Parse the output of an LLM
        Example:
            input_str = '''
                question: 今天天气怎么样？
                answer: 今天天气很好
                question: 中国首都？
                answer： 北京

                question:  上海被称为？
                answer： 魔都
            '''
            parser = QALinesParser()
            res = parser.parse(input_str)
            print(res)
            >> [{'question': '今天天气怎么样？', 'answer': '今天天气很好'},
            {'question': '中国首都？', 'answer': '北京'},
            {'question': '上海被称为？', 'answer': '魔都'}]

    """
    q_symbol: str = "question"
    a_symbol: str = "answer"

    def parse(self, text: str):
        """Parse the output to qa pairs."""
        # question_regex = re.compile(rf'{self.q_symbol}[:：]\s*(?P<que>.*?)$')
        # answer_regex = re.compile(rf'{self.a_symbol}[:：]\s*(?P<ans>.*?)$')
        question_regex = re.compile(rf'{self.q_symbol}[\d]*[:：]\s*(?P<que>.*?)\n')  # 多行匹配qa，需要匹配到每个换行符\n
        answer_regex = re.compile(rf'{self.a_symbol}[\d]*[:：]\s*(?P<ans>.*?)\n')
        m_q = re.finditer(question_regex, text)
        m_a = re.finditer(answer_regex, text)
        questions, answers = [], []
        if m_q and m_a:
            questions = [q['que'].strip() for q in m_q]
            answers = [a['ans'].strip() for a in m_a]
        # assert len(questions) == len(answers), "question and answer 's length mast same!"
        qa_pairs = []
        for que, ans in zip(questions, answers):
            qa_pairs.append({self.q_symbol: que, self.a_symbol: ans})
        return qa_pairs

    def get_format_instructions(self):
        format_str = f'''
        请以下面的格式输出：
        {self.q_symbol}: 生成的问题1\n
        {self.a_symbol}: 你生成的答案1\n
        {self.q_symbol}: 生成的问题2\n
        {self.a_symbol}: 你生成的答案2\n
        ...
        '''
        return format_str

# @register_parser("blocks")
class TextBlocksParser(DalchemyParser):
    """ 从文本中解析n个blocks ,注意：qa之间不能有空行
        Example:

        input_str = '''
            Q: 什么是人工智能？
            A: 人工智能（AI）是计算机科学的一个分支，致力于使计算机
            系统能够模仿人类智能的各种能力，如学习、推理、问题解决、感知、语言理解等。


            Q: 机器学习和深度学习有什么区别？
            A: 器学习是一种人工智能的子领域，它关注使计算机从数据中学习模式并做出预测或决策。
            而深度学习是机器学习的一种特定方法，它使用深层神经网络来进行学习和特征提取。
        '''
        from dalchemy.parsers import TextBlocksParser
        prefix_ls = ["Q:", "A:"]
        parser = TextBlocksParser(prefix_ls = prefix_ls)
        res = parser.parse(text)
        print(res[0])
        >> {'Q:': '什么是人工智能？', 'A:': "人工智能（AI）是计算机科学的一个分支，致力于使计算机系统能够模仿人类智能的各种能力，如学习、推理、问题解决、感知、语言理解等。"}

    """
    prefix_ls = ["Q", "A"]

    def parse(self, text: str):
        res = []
        # blocks = text.split("\n\n") # todo: 切分段落不够鲁棒 √
        blocks = txthelper.split_paragraphs_by_prefix(text, prefix=self.prefix_ls[0])
        for block in blocks:
            block = block.strip()
            if block == "": continue
            record = {}
            prefix_ls2 = self.prefix_ls[1:]
            for prefix_a, prefix_b in zip(self.prefix_ls, prefix_ls2):
                record[prefix_a] = txthelper.get_content_between_a_b(a=prefix_a, b=prefix_b, text=block)
            # final field
            last_prefix = self.prefix_ls[-1]
            record[last_prefix] = txthelper.get_content_after_b(b=last_prefix, text=block)
            res.append(record)
        return res

    def get_format_instructions(self) -> str:
        ''' 格式如下：块之间空两行，块内每个字段占多行
            Q: 文本块
            ...
            A: 文本块
            ...


            Q: 文本块
            ...
            A: 文本块
            ...
        '''
        block = [f"{prefix}: 文本块\n...\n" for prefix in self.prefix_ls] + ["\n\n"]
        format_str = "".join(block) * 2
        return format_str

